from baselines import BaseLine
from baselines.demo.HMM import HMM, TextSequence, viterbi
import pickle
import numpy as np
import os

basedir = os.path.dirname(os.path.abspath(__file__))
dictname = os.path.join(basedir, "./dict.txt")
dictname2 = os.path.join(basedir, "./pos-tag.txt")
modelname = os.path.join(basedir, "./model.hmm")

# print("正在加载字典...")
idx2tag = np.loadtxt(dictname, dtype=str)
# print("字典加载完毕")
# print("正在加载字典2...")
tag2zh = np.loadtxt(dictname2, dtype=str, encoding="utf-8")
tag2zhDict = {}
for i in range(tag2zh.shape[0]):
    tag2zhDict[tag2zh[i, 0]] = tag2zh[i, 1]


def tag2pure_fn(tag:str)->str:
    tag_pure =tag.split("-")[1]
    return tag_pure

# print("字典加载完毕2")
# print(tag2zhDict)
# print("正在重新加载模型...")
with open(file=modelname, mode="br") as file:
    hmm2 = pickle.load(file)
# print("模型加载成功")
# print("开始预测...")

class DemoBaseLine(BaseLine):
    """
    基于HMM实现的样例基线模型
    """
    def preidct(self, sentence:str):
        chars = list(sentence)
        states = viterbi(hmm2, hmm2.getIndex(chars))
        results = []
        tags = [idx2tag[idx] for idx in states]
        word = []
        currentTag = 'B-n'
        for i in range(len(tags)):
            tag = tags[i]
            if(tag != 'I'):
                pure_tag = tag2pure_fn(currentTag)
                results.append((''.join(word), pure_tag, tag2zhDict[pure_tag] ))
                word = []
                currentTag = tag
            word.append(chars[i])

        pure_tag = tag2pure_fn(currentTag)
        results.append((''.join(word), pure_tag, tag2zhDict[pure_tag] ))
        results = results[1:]
        return results


if __name__ == '__main__':
    demo = DemoBaseLine()
    print("结果：")
    print(demo.preidct("嘉然今天吃什么？"))
    print(demo.preidct("如果有一天我突然死了，你会记得我吗？"))
    print(demo.preidct("中华人民共和国今天成立了！"))
    print(demo.preidct("新华社北京12月14日电 12月14日，国家主席习近平就美国肯塔基、伊利诺伊、阿肯色、密西西比、田纳西、密苏里等中部6州遭遇多场龙卷风袭击并造成重大人员伤亡和财产损失，向美国总统拜登致慰问电，代表中国政府和人民向受灾的美国人民表示深切同情和诚挚慰问。"))